import pandas as pd
import numpy as np
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sklearn.metrics import pairwise_distances
# from sklearn.neighbors import NearestNeighbors
from scipy.spatial.distance import euclidean, mahalanobis,cdist
from scipy.stats import wasserstein_distance
import seaborn as sns
import matplotlib.pyplot as plt
from random_case import find_cases
from tqdm import tqdm
# import numpy as np
from sklearn.manifold import Isomap
from datetime import datetime
import os
import gc
from collections import defaultdict
class Similarity():
    def __init__(self,file_path) -> None:
        self.file_dir = file_path
        self.cases = find_cases(file_path)
        self.Project = self.cases[0][0].split('\\')[0]
        self.cases_sample_type = pd.DataFrame(self.cases).iloc[0:,1].values
        self.temp_df = pd.DataFrame()
        self.is_df_exist = False
        self.best_k = {
            'TCGA-BLCA':3,
            'TCGA-BRCA':3,
            'TCGA-LGG':4,
            'TCGA-LUAD':4,
            'TCGA-LUSC':4,
        }
    def pre_process(self):
        self.temp_df = pd.read_csv(self.file_dir+'.tsv',sep='\t',index_col=0,header=0).transpose()
        self.temp_df2 = pd.read_csv(self.file_dir+'_2D.tsv',sep='\t',index_col=0,header=0)
        print(self.temp_df,self.temp_df.shape)
        
        '''调试用, 绘制不同k值时的曲线，寻找拐点，不用时注释'''
        # wcss = []
        # for k in range(1,11):
        #     kmeans = KMeans(n_clusters=k, random_state=0,n_init='auto')
        #     clusters = kmeans.fit_predict(self.temp_df2)
        #     wcss.append(kmeans.inertia_)
        # # 绘制WCSS与K值的关系图
        # plt.figure(figsize=(10,8))
        # plt.plot(range(1, 11), wcss)
        # plt.title(f'The Elbow Method showing the optimal k -- {self.Project}')
        # plt.xlabel('Number of clusters')
        # plt.ylabel('WCSS')

        # current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # filename = f'The Elbow Method showing the optimal k -- {self.Project}-{current_time}.png'
        # output_folder = f'k_line'
        # if not os.path.exists(output_folder):
        #     os.makedirs(output_folder)
        # plt.savefig(os.path.join(output_folder, filename))
        # plt.show()
        '''调试用结束'''
        
        self.k = self.best_k[self.file_dir]
        kmeans = KMeans(n_clusters=self.k, random_state=0,n_init='auto')
        clusters = kmeans.fit_predict(self.temp_df2)
        # 将聚类的结果用作排序
        self.sort_indices = np.argsort(clusters)
        self.tsne = TSNE(n_components=64, random_state=0,method='exact')
        np_reduced = self.tsne.fit_transform(self.temp_df)
        self.sorted_np_reduced = np_reduced[self.sort_indices]
        # 使用主成分分析
        self.pca = PCA(n_components=64)
        np_reduced_pca = self.pca.fit_transform(self.temp_df)
        self.sorted_np_reduced_pca = np_reduced_pca[self.sort_indices]

        self.sorted_np_origin = self.temp_df.values[self.sort_indices]
        pass

    def draw_pearson_similarity_matrix(self,reduced_enable = True,embedding = 'tsne'):
        if embedding == 'tsne':
            if reduced_enable:
                sorted_df = pd.DataFrame(self.sorted_np_reduced)
                distance_file_name = 'pearson_distance_tsne.tsv'
            else:
                sorted_df = pd.DataFrame(self.sorted_np_origin)
                distance_file_name = 'pearson_distance_raw.tsv'
        elif embedding == 'pca':
            if reduced_enable:
                sorted_df = pd.DataFrame(self.sorted_np_reduced_pca)
                distance_file_name = 'pearson_distance_pca.tsv'
            else:
                sorted_df = pd.DataFrame(self.sorted_np_origin)
                distance_file_name = 'pearson_distance_raw.tsv'
        else:
            ValueError('embedding type error')
        # sorted_df = sorted_df.transpose()
        # pearman_matrix = sorted_df.corr(method='pearson').apply(lambda x : 1.0 - x)
        pearman_matrix = cdist(sorted_df,sorted_df,metric='correlation')
        print(pearman_matrix,pearman_matrix.shape)
        plt.figure(figsize=(10, 8))
        sns.heatmap(pearman_matrix, cmap='coolwarm', annot=False, fmt='.2f',xticklabels=False,yticklabels=False)
        plt.title('Similarity Matrix - pearson')
        # 获取当前的日期和时间
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # 构建文件名，例如 '2024-04-13-15-30-45.png'
        filename = f'pearson-embedding-{embedding}-reduced-{reduced_enable}-{current_time}.png'
        output_folder = f'embedding-pic\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))

        # 构建距离矩阵文件名'
        output_folder = f'distance-matrix\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        df_pearson = pd.DataFrame(pearman_matrix)
        df_pearson.to_csv(os.path.join(output_folder, distance_file_name))
        plt.show()
        return
        
    def draw_Euclidean_similarity_matrix(self,reduced_enable = True,embedding = 'tsne'):
        '''
            函数简介：绘制欧式距离
        '''
        if embedding == 'tsne':
            if reduced_enable:
                euclidean_matrix = pairwise_distances(self.sorted_np_reduced)
                distance_file_name = 'euclidean_distance_tsne.tsv'
            else:
                euclidean_matrix = pairwise_distances(self.sorted_np_origin)
                distance_file_name = 'euclidean_distance_raw.tsv'
        elif embedding == 'pca':
            if reduced_enable:
                euclidean_matrix = pairwise_distances(self.sorted_np_reduced_pca)
                distance_file_name = 'euclidean_distance_pca.tsv'
            else:
                euclidean_matrix = pairwise_distances(self.sorted_np_origin)
                distance_file_name = 'euclidean_distance_raw.tsv'
        else:
            ValueError('embedding type error')
        
        print(euclidean_matrix,type(euclidean_matrix))
        plt.figure(figsize=(10, 8))
        sns.heatmap(euclidean_matrix, cmap='coolwarm', annot=False, fmt='.2f',xticklabels=False,yticklabels=False)
        plt.title('Similarity Matrix - Euclidean Distance')
         # 获取当前的日期和时间
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # 构建文件名，例如 '2024-04-13-15-30-45.png'
        filename = f'Euclidean-embedding-{embedding}-reduced-{reduced_enable}-{current_time}.png'
        output_folder = f'embedding-pic\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))

        output_folder = f'distance-matrix\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        df_pearson = pd.DataFrame(euclidean_matrix)
        df_pearson.to_csv(os.path.join(output_folder, distance_file_name))
        plt.show()
        pass
    
    def draw_Mahalanobis_similarity_matrix(self,reduced_enable = True,embedding = 'tsne'):
        '''
            函数简介：绘制马氏距离
        '''
        if embedding == 'tsne':
            if reduced_enable:
                Mahalanobis_matrix = pairwise_distances(self.sorted_np_reduced,metric='mahalanobis')
                distance_file_name = 'mahalanobis_distance_tsne.tsv'
            else:
                Mahalanobis_matrix = pairwise_distances(self.sorted_np_origin[0:15,:],metric='mahalanobis')
                # Mahalanobis_matrix = cdist(self.sorted_np_origin[0:15,:], self.sorted_np_origin[0:15,:], metric='mahalanobis')
                distance_file_name = 'mahalanobis_distance_raw.tsv'
        elif embedding == 'pca':
            if reduced_enable:
                Mahalanobis_matrix = pairwise_distances(self.sorted_np_reduced_pca,metric='mahalanobis')
                distance_file_name = 'mahalanobis_distance_pca.tsv'
            else:
                Mahalanobis_matrix = pairwise_distances(self.sorted_np_origin,metric='mahalanobis')
                distance_file_name = 'mahalanobis_distance_raw.tsv'
        else:
            ValueError('embedding type error')

        # print(Mahalanobis_matrix,type(Mahalanobis_matrix))
        plt.figure(figsize=(10, 8))
        sns.heatmap(Mahalanobis_matrix, cmap='coolwarm', annot=False, fmt='.2f',xticklabels=False,yticklabels=False)
        plt.title('Similarity Matrix - Mahalanobis Distance')
         # 获取当前的日期和时间
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # 构建文件名，例如 '2024-04-13-15-30-45.png'
        filename = f'Mahalanobis-embedding-{embedding}-reduced-{reduced_enable}-{current_time}.png'
        output_folder = f'embedding-pic\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))

        output_folder = f'distance-matrix\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        df_pearson = pd.DataFrame(Mahalanobis_matrix)
        df_pearson.to_csv(os.path.join(output_folder, distance_file_name))
        plt.show()
        pass

    def draw_cosine_similarity_matrix(self,reduced_enable = True,embedding = 'tsne'):
        # 余弦距离
        if embedding == 'tsne':
            if reduced_enable:
                cosine_sim = pairwise_distances(self.sorted_np_reduced,metric='cosine')
                distance_file_name = 'cosine_distance_tsne.tsv'
            else:
                cosine_sim = pairwise_distances(self.sorted_np_origin,metric='cosine')
                distance_file_name = 'cosine_distance_raw.tsv'
        elif embedding == 'pca':
            if reduced_enable:
                cosine_sim = pairwise_distances(self.sorted_np_reduced_pca,metric='cosine')
                distance_file_name = 'cosine_distance_pca.tsv'
            else:
                cosine_sim = pairwise_distances(self.sorted_np_origin,metric='cosine')
                distance_file_name = 'cosine_distance_raw.tsv'
        else:
            ValueError('embedding type error')
        # cosine_sim = pairwise_distances(self.sorted_np_reduced,metric='cosine')
        # print(cosine_sim,type(cosine_sim))
        plt.figure(figsize=(10, 8))
        sns.heatmap(cosine_sim, cmap='coolwarm', annot=False, fmt='.2f',xticklabels=False,yticklabels=False)
        plt.title('Similarity Matrix - cosine Distance')
         # 获取当前的日期和时间
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # 构建文件名，例如 '2024-04-13-15-30-45.png'
        filename = f'cosine-embedding-{embedding}-reduced-{reduced_enable}-{current_time}.png'
        output_folder = f'embedding-pic\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))

        output_folder = f'distance-matrix\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        df_pearson = pd.DataFrame(cosine_sim)
        df_pearson.to_csv(os.path.join(output_folder, distance_file_name))
        plt.show()

    def draw_geodesic_similarity_matrix(self,reduced_enable = True,embedding = 'tsne'):
        if reduced_enable:
            if embedding == 'tsne':
                isomap = Isomap(n_components=2,n_neighbors=5,path_method="auto")
                data_2d = isomap.fit_transform(X=self.sorted_np_reduced)
                distance_file_name = 'geodesic_distance_tsne.tsv'
            elif embedding == 'pca':
                isomap = Isomap(n_components=2,n_neighbors=5,path_method="auto")
                data_2d = isomap.fit_transform(X=self.sorted_np_reduced_pca)
                distance_file_name = 'geodesic_distance_pca.tsv'
        else:
            isomap = Isomap(n_components=2,n_neighbors=5,path_method="auto")
            data_2d = isomap.fit_transform(X=self.temp_df.values[self.sort_indices])
            distance_file_name = 'geodesic_distance_raw.tsv'
        
        geo_distance_matrix = isomap.dist_matrix_ # 测地距离矩阵，shape=[n_sample,n_sample]
        print(geo_distance_matrix,type(geo_distance_matrix),geo_distance_matrix.shape)
        plt.figure(figsize=(10, 8))
        sns.heatmap(geo_distance_matrix, cmap='coolwarm', annot=False, fmt='.2f',xticklabels=False,yticklabels=False)
        plt.title('Similarity Matrix - Geodesic Distance')
         # 获取当前的日期和时间
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # 构建文件名，例如 '2024-04-13-15-30-45.png'
        filename = f'Geodesic-embedding-{embedding}-reduced-{reduced_enable}-{current_time}.png'
        output_folder = f'embedding-pic\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))

        output_folder = f'distance-matrix\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        df_pearson = pd.DataFrame(geo_distance_matrix)
        df_pearson.to_csv(os.path.join(output_folder, distance_file_name))
        plt.show() 

    def draw_wasserstein_similarity_matrix(self,reduced_enable = True,embedding = 'tsne'):
        if embedding == 'tsne':
            if reduced_enable:
                sorted_df = pd.DataFrame(self.sorted_np_reduced)
                wasserstein_matrix = cdist(sorted_df,sorted_df, lambda u, v: wasserstein_distance(u,v))
                distance_file_name = 'wasserstein_distance_tsne.tsv'
            else:
                sorted_df = pd.DataFrame(self.sorted_np_origin).iloc[0:15,:]
                wasserstein_matrix = cdist(sorted_df,sorted_df, lambda u, v: wasserstein_distance(u,v))
                distance_file_name = 'wasserstein_distance_raw.tsv'
        elif embedding == 'pca':
            if reduced_enable:
                sorted_df = pd.DataFrame(self.sorted_np_reduced_pca)
                wasserstein_matrix = cdist(sorted_df,sorted_df, lambda u, v: wasserstein_distance(u,v))
                distance_file_name = 'wasserstein_distance_pca.tsv'
            else:
                sorted_df = pd.DataFrame(self.sorted_np_origin)
                wasserstein_matrix = cdist(sorted_df,sorted_df, lambda u, v: wasserstein_distance(u,v))
                distance_file_name = 'wasserstein_distance_raw.tsv'
        else:
            ValueError('embedding error type')
                
        # print(wasserstein_matrix,type(wasserstein_matrix),wasserstein_matrix.shape)
        plt.figure(figsize=(10, 8))
        sns.heatmap(wasserstein_matrix, cmap='coolwarm', annot=False, fmt='.2f',xticklabels=False,yticklabels=False)
        plt.title('Similarity Matrix - wasserstein Distance')
         # 获取当前的日期和时间
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        # 构建文件名，例如 '2024-04-13-15-30-45.png'
        filename = f'wasserstein-embedding-{embedding}-reduced-{reduced_enable}-{current_time}.png'
        output_folder = f'embedding-pic\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))
        plt.show() 

        output_folder = f'distance-matrix\\{self.Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        df_pearson = pd.DataFrame(wasserstein_matrix)
        df_pearson.to_csv(os.path.join(output_folder, distance_file_name))
        pass

if __name__ == '__main__':

    # draw_map = Similarity('TCGA-BLCA') 
    file_list=[
        'TCGA-BLCA',
        'TCGA-BRCA',
        'TCGA-LGG',
        'TCGA-LUAD',
        'TCGA-LUSC',
    ]
    for file_dir in file_list:
        draw_map = Similarity(file_dir) 
        draw_map.pre_process()

        #降至64维，tsne非线性降维
        draw_map.draw_pearson_similarity_matrix(reduced_enable=True,embedding='tsne')
        draw_map.draw_Euclidean_similarity_matrix(reduced_enable=True,embedding='tsne')
        draw_map.draw_Mahalanobis_similarity_matrix(reduced_enable=True,embedding='tsne')
        draw_map.draw_cosine_similarity_matrix(reduced_enable=True,embedding='tsne')
        draw_map.draw_geodesic_similarity_matrix(reduced_enable=True,embedding='tsne')
        draw_map.draw_wasserstein_similarity_matrix(reduced_enable=True,embedding='tsne')
        #降至64维，主成分分析法
        draw_map.draw_pearson_similarity_matrix(reduced_enable=True,embedding='pca')
        draw_map.draw_Euclidean_similarity_matrix(reduced_enable=True,embedding='pca')
        draw_map.draw_Mahalanobis_similarity_matrix(reduced_enable=True,embedding='pca')
        draw_map.draw_cosine_similarity_matrix(reduced_enable=True,embedding='pca')
        draw_map.draw_geodesic_similarity_matrix(reduced_enable=True,embedding='pca')
        draw_map.draw_wasserstein_similarity_matrix(reduced_enable=True,embedding='pca')
        gc.collect()
        #不降维，使用原始数据
        draw_map.draw_pearson_similarity_matrix(reduced_enable=False)
        draw_map.draw_Euclidean_similarity_matrix(reduced_enable=False)
        draw_map.draw_Mahalanobis_similarity_matrix(reduced_enable=False)
        draw_map.draw_cosine_similarity_matrix(reduced_enable=False)
        draw_map.draw_geodesic_similarity_matrix(reduced_enable=False)
        draw_map.draw_wasserstein_similarity_matrix(reduced_enable=False)
        gc.collect()
        